package com.hc.summary.asm;

import org.junit.Test;
import org.objectweb.asm.AnnotationVisitor;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.AdviceAdapter;
import org.objectweb.asm.commons.Method;

import java.io.FileInputStream;
import java.io.FileOutputStream;

public class ASMTool {
    @Test
    public void test(){
        try {
            FileInputStream fis = new FileInputStream("E:\\ASProject\\Summary\\app\\src\\test\\ASMBean.class");//待修改的类
            ClassReader classReader = new ClassReader(fis);
            ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
            classReader.accept(new CustomClassVisitor(Opcodes.ASM7,classWriter),ClassReader.EXPAND_FRAMES);
            byte[] bytes = classWriter.toByteArray();
            FileOutputStream fos = new FileOutputStream("E:\\ASProject\\Summary\\app\\src\\test\\ASMBean2.class");//修改后的类
            fos.write(bytes);
            fos.close();
            fis.close();
        }catch (Exception e){
            e.printStackTrace();
        }
    }

    private static class CustomClassVisitor extends ClassVisitor{
        public CustomClassVisitor(int api, ClassVisitor classVisitor) {
            super(api, classVisitor);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            MethodVisitor methodVisitor = super.visitMethod(access, name, descriptor, signature, exceptions);
            return new CustomMethodVisitor(api,methodVisitor,access,name,descriptor);
        }
    }

    private static class CustomMethodVisitor extends AdviceAdapter {
        boolean inject = false;
        int startIndex;
        int endIndex;
        public CustomMethodVisitor(int api, MethodVisitor methodVisitor,int access,String name,String descriptor) {
            super(api, methodVisitor,access,name,descriptor);
        }

        @Override
        public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
            if("Lcom/hc/summary/asm/ASMPoint;".equals(descriptor)){
                inject = true;
            }else {
                inject = false;
            }
            return super.visitAnnotation(descriptor, visible);
        }

        @Override
        protected void onMethodEnter() {
            super.onMethodEnter();
            if(!inject){
                return;
            }
            invokeStatic(Type.getType("Ljava/lang/System;"),new Method("currentTimeMillis","()J"));
            startIndex = newLocal(Type.LONG_TYPE);
            storeLocal(startIndex);
        }

        @Override
        protected void onMethodExit(int opcode) {
            super.onMethodExit(opcode);
            if(!inject){
                return;
            }
            invokeStatic(Type.getType("Ljava/lang/System;"),new Method("currentTimeMillis","()J"));
            endIndex = newLocal(Type.LONG_TYPE);
            storeLocal(endIndex);
            getStatic(Type.getType("Ljava/lang/System;"),"out",Type.getType("Ljava/io/PrintStream;"));
            newInstance(Type.getType("Ljava/lang/StringBuilder;"));
            dup();
            invokeConstructor(Type.getType("Ljava/lang/StringBuilder;"),new Method("<init>","()V"));
            visitLdcInsn("execute time=");
            invokeVirtual(Type.getType("Ljava/lang/StringBuilder;"),new Method("append","(Ljava/lang/String;)Ljava/lang/StringBuilder;"));
            loadLocal(endIndex);
            loadLocal(startIndex);
            math(SUB,Type.LONG_TYPE);
            invokeVirtual(Type.getType("Ljava/lang/StringBuilder;"),new Method("append","(J)Ljava/lang/StringBuilder;"));
            visitLdcInsn(" ms");
            invokeVirtual(Type.getType("Ljava/lang/StringBuilder;"),new Method("append","(Ljava/lang/String;)Ljava/lang/StringBuilder;"));
            invokeVirtual(Type.getType("Ljava/lang/StringBuilder;"),new Method("toString","()Ljava/lang/String;"));
            invokeVirtual(Type.getType("Ljava/io/PrintStream;"),new Method("println","(Ljava/lang/String;)V;"));
        }
    }
}
